import sys
import sqlite3
import math
import ast
import json
from PySide6.QtCore import (Qt, QAbstractItemModel, QModelIndex, QEvent)
from PySide6.QtWidgets import (QApplication, QWidget, QVBoxLayout, QHBoxLayout, QSplitter, QFormLayout, QLabel, QMessageBox,QFrame)
from PySide6.QtGui import QClipboard
from qfluentwidgets import (LineEdit,TreeView,PushButton,PrimaryPushButton,InfoBarPosition,InfoBar,ComboBox,TextEdit)

# 定义一个安全的求值函数，只允许使用预定义的安全函数和操作符
def safe_eval(formula, variables):
    # 只允许特定的数学函数和常量
    allowed_math_functions = {
        'sqrt': math.sqrt,
        'sin': math.sin,
        'cos': math.cos,
        'tan': math.tan,
        'exp': math.exp,
        'log': math.log,
        'pi': math.pi,
        'e': math.e,
        # 你可以根据需要添加更多的数学函数和常量
    }

    # 创建一个受限的 math 模块对象
    allowed_math_module = type('math', (object,), allowed_math_functions)()

    allowed_names = {
        **variables,
        'math': allowed_math_module  # 使用受限的 math 模块对象
    }

    try:
        # Parse the formula into an AST
        parsed_ast = ast.parse(formula, mode='eval')
        
        # Walk through the AST and check for disallowed names or calls
        for node in ast.walk(parsed_ast):
            if isinstance(node, ast.Name):
                if node.id not in allowed_names:
                    raise SyntaxError(f"不允许使用的变量或函数: {node.id}")
            elif isinstance(node, ast.Call):
                func = node.func
                if isinstance(func, ast.Attribute):
                    # Check if the attribute is allowed (e.g., math.sqrt)
                    if not (isinstance(func.value, ast.Name) and func.value.id == 'math' and hasattr(allowed_math_module, func.attr)):
                        raise SyntaxError(f"不允许的函数调用: {ast.dump(func)}")
                elif isinstance(func, ast.Name) and not hasattr(allowed_math_module, func.id) and func.id not in variables:
                    raise SyntaxError(f"不允许的函数调用: {ast.dump(func)}")
        
        # Compile the AST to a code object
        compiled_code = compile(parsed_ast, '<string>', 'eval')
        
        # Evaluate the compiled code with the restricted environment
        result = eval(compiled_code, {"__builtins__": None}, allowed_names)
        return result, ""  # 返回结果和空字符串作为错误信息
    except ValueError as e:
        # 捕获数学域错误并返回有意义的错误信息
        return None, f"数学域错误: {str(e)}"
    except Exception as e:
        return None, str(e)  # 返回 None 和其他类型的错误信息

class FormulaModel(QAbstractItemModel):
    def __init__(self):
        super().__init__()
        self.load_data()

    def load_data(self):
        conn = sqlite3.connect('formulas.db')
        c = conn.cursor()
        
        # 加载分类
        c.execute('SELECT id, name FROM categories')
        self.categories = [{'id': row[0], 'name': row[1], 'children': []} for row in c.fetchall()]
        
        # 加载公式并将其分配到对应的分类
        c.execute('SELECT id, name, formula, variables, category_id FROM formulas')
        formulas = [{'id': row[0], 'name': row[1], 'value': row[2], 'variables': ast.literal_eval(row[3]), 'category_id': row[4]} for row in c.fetchall()]
        for formula in formulas:
            for category in self.categories:
                if category['id'] == formula['category_id']:
                    category['children'].append(formula)
                    break
        
        conn.close()

    def index(self, row, column, parent=QModelIndex()):
        if not self.hasIndex(row, column, parent):
            return QModelIndex()
        
        if not parent.isValid():  # 如果没有父级，表示是分类
            if 0 <= row < len(self.categories):
                return self.createIndex(row, column, self.categories[row])
        else:  # 否则，表示是分类下的公式
            parent_item = parent.internalPointer()
            if 0 <= row < len(parent_item['children']):
                return self.createIndex(row, column, parent_item['children'][row])
        
        return QModelIndex()

    def parent(self, index):
        if not index.isValid():
            return QModelIndex()
        
        item = index.internalPointer()
        if 'category_id' in item:  # 如果有 category_id，则表示是公式，返回其分类的索引
            for i, category in enumerate(self.categories):
                if category['id'] == item['category_id']:
                    return self.createIndex(i, 0, category)
        return QModelIndex()  # 否则是分类，返回无效索引

    def rowCount(self, parent=QModelIndex()):
        if not parent.isValid():
            return len(self.categories)  # 分类的数量
        else:
            parent_item = parent.internalPointer()
            return len(parent_item.get('children', []))  # 公式数量

    def columnCount(self, parent=QModelIndex()):
        return 2  # 名称和值两列
    
    def headerData(self, section, orientation, role=Qt.DisplayRole):
        if orientation == Qt.Horizontal and role == Qt.DisplayRole:
            if section == 0:
                return "名称"
            elif section == 1:
                return "公式"
        return None

    def data(self, index, role=Qt.DisplayRole):
        if not index.isValid():
            return None
        
        item = index.internalPointer()
        if role == Qt.DisplayRole:
            if index.column() == 0:
                return item['name']
            elif index.column() == 1:
                # return item.get('value', '')  # 对于分类，这一列为空
                if 'value' in item:  # 只有公式有 'value' 键，分类没有
                    return item['value']
        return None

    def get_formula_by_index(self, index):
        if not index.isValid():
            return None
        item = index.internalPointer()
        if 'category_id' in item:  # 确保我们获取的是公式，而不是分类
            return item
        return None
    
    def save_or_update_formula(self, formula):
        # 连接到数据库
        conn = sqlite3.connect('formulas.db')
        cursor = conn.cursor()
        try:
            if 'id' in formula:
                # 更新现有公式
                cursor.execute('''
                    UPDATE formulas
                    SET name = ?, category_id = ?, formula = ?, variables = ?
                    WHERE id = ?
                ''', (formula['name'], formula['category_id'], formula['formula'], formula['variables'], formula['id']))
            else:
                # 插入新公式
                cursor.execute('''
                    INSERT INTO formulas (name, category_id, formula, variables)
                    VALUES (?, ?, ?, ?)
                ''', (formula['name'], formula['category_id'], formula['formula'], formula['variables']))

            # 提交更改
            conn.commit() 
            # 刷新数据
            # self.load_data()
            return True, "保存成功"
        except sqlite3.Error as e:
            print(f"Error: {e}")
            return False, f"保存失败: {e}"
        finally:
            # 关闭连接
            conn.close()

    def delete_formula(self, formula_id):
        # 连接到数据库
        conn = sqlite3.connect('formulas.db')
        cursor = conn.cursor()
        try:
            # 删除公式
            cursor.execute('DELETE FROM formulas WHERE id =?', (formula_id,))
            # 提交更改
            conn.commit()
            return True, "删除成功"
        except sqlite3.Error as e:
            print(f"Error: {e}")
            return False, f"删除失败: {e}"
        finally:
            # 关闭连接
            conn.close()
        


# 计算器主窗口类
class Calculator(QFrame):
    def __init__(self, text: str, parent=None):
        super().__init__(parent=parent)
        self.setObjectName(text.replace(' ', '-'))

        main_layout = QHBoxLayout()
        splitter = QSplitter(Qt.Horizontal)

        # 左侧：公式选择树
        # 右侧：公式表达式输入框、变量输入表单、表达式显示框、结果显示框和公式信息标签
        left_widget = QWidget()
        left_layout = QVBoxLayout()
        self.tree_view = TreeView()
        self.model = FormulaModel()
        self.tree_view.setModel(self.model)
        self.tree_view.clicked.connect(self.on_tree_clicked)
        # 为 TreeView 添加边框，同时保持原有样式
        self.update_tree_view_border()
        self.add_edit_button = PushButton("新增/编辑")
        self.add_edit_button.clicked.connect(self.on_add_edit_clicked)
        left_layout.addWidget(self.tree_view)
        left_layout.addWidget(self.add_edit_button)
        left_widget.setLayout(left_layout)
        splitter.addWidget(left_widget)

        # 右侧：公式表达式输入框、变量输入表单、表达式显示框、结果显示框和公式信息标签
        self.right_widget = QWidget()
        right_layout = QVBoxLayout()

        self.input_edit = LineEdit()
        self.input_edit.setReadOnly(True)
        self.input_edit.setPlaceholderText("选择公式后显示表达式")
        right_layout.addWidget(self.input_edit)

        self.variables_layout = QFormLayout()
        right_layout.addLayout(self.variables_layout)

        # 表达式显示框及复制按钮
        expression_layout = QHBoxLayout()
        self.expression_edit = LineEdit()
        self.expression_edit.setReadOnly(True)
        self.expression_edit.setPlaceholderText("带入变量值的表达式")
        copy_expression_button = PushButton('复制表达式')
        copy_expression_button.clicked.connect(self.copy_expression)
        expression_layout.addWidget(self.expression_edit)
        expression_layout.addWidget(copy_expression_button)
        right_layout.addLayout(expression_layout)

        # 结果显示框及复制按钮
        result_layout = QHBoxLayout()
        self.result_edit = LineEdit()
        self.result_edit.setPlaceholderText("计算结果")
        self.result_edit.setReadOnly(True)
        copy_result_button = PushButton('复制结果')
        copy_result_button.clicked.connect(self.copy_result)
        result_layout.addWidget(self.result_edit)
        result_layout.addWidget(copy_result_button)
        right_layout.addLayout(result_layout)

        calculate_button = PrimaryPushButton('计算（Enter）')
        calculate_button.clicked.connect(self.calculate)
        right_layout.addWidget(calculate_button)

        self.right_widget.setLayout(right_layout)
        # 隐藏的公式编辑窗口
        self.formula_widget = QWidget()
        self.formula_layout = QVBoxLayout()
        # 公式名称输入框
        self.formula_name_edit = LineEdit()
        self.formula_name_edit.setPlaceholderText("公式名称")
        self.formula_layout.addWidget(self.formula_name_edit)

        # 公式分类选择框
        self.formula_category_combo_box = ComboBox()
        self.formula_category_combo_box.addItems([category['name'] for category in self.model.categories])
        for key, value in enumerate(self.model.categories):
            self.formula_category_combo_box.setItemData(key, value['id'])
        self.formula_layout.addWidget(self.formula_category_combo_box)

        # 公式表达式输入框
        self.formula_expression_edit = LineEdit()
        self.formula_expression_edit.setPlaceholderText("公式表达式")
        self.formula_layout.addWidget(self.formula_expression_edit)

        # 变量输入框
        self.variables_textedit = TextEdit()
        self.variables_textedit.setPlaceholderText("变量列表json格式")
        self.formula_layout.addWidget(self.variables_textedit)
        
        # 加弹簧
        self.formula_layout.addStretch()

        # 保存按钮
        self.save_return_layout = QHBoxLayout()
        self.save_button = PrimaryPushButton('保存')
        self.save_button.clicked.connect(lambda:self.save_formula(self.current_formula_id))

        self.return_button = PushButton('返回')
        self.return_button.clicked.connect(self.on_back_button_clicked)
        self.save_return_layout.addWidget(self.save_button)
        self.save_return_layout.addWidget(self.return_button)
        # 保存按钮占比大
        self.return_button.setFixedWidth(80)
        self.formula_layout.addLayout(self.save_return_layout)

        self.del_button = PushButton('删除')
        self.del_button.clicked.connect(self.on_delete_button_clicked)
        self.formula_layout.addWidget(self.del_button)
        self.formula_widget.setLayout(self.formula_layout)

        splitter.addWidget(self.right_widget)
        splitter.addWidget(self.formula_widget)
        self.formula_widget.hide()

        splitter.setStretchFactor(1, 1)

        main_layout.addWidget(splitter)
        self.setLayout(main_layout)

        self.current_formula = None

        # 添加事件过滤器以捕获回车键事件
        self.input_edit.installEventFilter(self)
        self.expression_edit.installEventFilter(self)
        self.result_edit.installEventFilter(self)
        self.installEventFilter(self)

        #当前公式id
        self.current_formula_id = None

    def eventFilter(self, obj, event):
        if event.type() == QEvent.KeyPress and (event.key() == Qt.Key_Return or event.key() == Qt.Key_Enter):
            if obj == self.input_edit:
                # 如果焦点在公式表达式输入框，跳到第一个变量输入框
                if self.variables_layout.rowCount() > 0:
                    first_var_edit = self.variables_layout.itemAt(0, QFormLayout.FieldRole).widget()
                    first_var_edit.setFocus()
            elif obj == self.expression_edit:
                # 如果焦点在表达式显示框，跳到第一个变量输入框
                if self.variables_layout.rowCount() > 0:
                    first_var_edit = self.variables_layout.itemAt(0, QFormLayout.FieldRole).widget()
                    first_var_edit.setFocus()
            elif obj == self.result_edit:
                # 如果焦点在结果显示框，执行计算
                self.calculate()
            else:
                # 如果焦点在变量输入框，跳到下一个变量输入框或执行计算
                for i in range(self.variables_layout.rowCount()):
                    var_edit = self.variables_layout.itemAt(i, QFormLayout.FieldRole).widget()
                    if var_edit == obj:
                        if i < self.variables_layout.rowCount() - 1:
                            next_var_edit = self.variables_layout.itemAt(i + 1, QFormLayout.FieldRole).widget()
                            next_var_edit.setFocus()
                        else:
                            # 如果是最后一个变量输入框，执行计算
                            self.calculate()
                        break
            return True
        if event.type() == QEvent.MouseButtonPress and obj is self:
            if not self.tree_view.rect().contains(event.pos()):
                self.tree_view.clearSelection()
                self.tree_view.selectionModel().clear()
                
        return super().eventFilter(obj, event)

    def update_tree_view_border(self):
        # 获取当前样式表
        current_style = self.tree_view.styleSheet()  
        
        # 定义新的边框样式
        new_style = """
            QTreeView {
                border: 1px solid #EAEAEA; /* 深色模式下的边框颜色 */
                border-radius: 2px;
            }
        """
        
        # 合并样式
        combined_style = current_style + new_style if current_style else new_style
        self.tree_view.setStyleSheet(combined_style)


    def copy_expression(self):
        """Copy the expression to the clipboard."""
        if not self.expression_edit.text():
            InfoBar.error(
                    title='提示',
                    content='没有可复制的表达式',
                    parent=self,
                    position=InfoBarPosition.TOP,
                    duration=2000
                )
            return
        clipboard = QApplication.clipboard()
        clipboard.setText(self.expression_edit.text())
        InfoBar.success(
                    title='成功',
                    content='表达式已复制到剪贴板',
                    parent=self,
                    position=InfoBarPosition.TOP,
                    duration=2000
        )

    def copy_result(self):
        """Copy the result to the clipboard."""
        if not self.result_edit.text():
            InfoBar.error(
                    title='提示',
                    content='没有可复制的结果',
                    parent=self,
                    position=InfoBarPosition.TOP,
                    duration=2000
                )
            return
        clipboard = QApplication.clipboard()
        clipboard.setText(self.result_edit.text())
        InfoBar.success(
                    title='成功',
                    content='结果已复制到剪贴板',
                    parent=self,
                    position=InfoBarPosition.TOP,
                    duration=2000
        )

    def on_tree_clicked(self, index):
        self.current_formula = self.model.get_formula_by_index(index)
        if self.current_formula:
            self.input_edit.setText(self.current_formula['value'])

            # Clear previous variable inputs
            while self.variables_layout.rowCount() > 0:
                self.variables_layout.removeRow(0)

            # Add new variable inputs with default values of 1
            for var, comment in self.current_formula['variables'].items():
                var_edit = LineEdit()  # Default value
                var_edit.setText("1")
                comment_label = QLabel(var+"("+comment +")")
                self.variables_layout.addRow(comment_label, var_edit)
                # 添加事件过滤器到每个变量输入框
                var_edit.installEventFilter(self)

            self.result_edit.clear()
            self.expression_edit.clear()

        else:  # 如果选择了分类而不是公式
            self.input_edit.clear()
            while self.variables_layout.rowCount() > 0:
                self.variables_layout.removeRow(0)
            self.result_edit.clear()
            self.expression_edit.clear()
            # self.formula_info_label.setText("")

    def calculate(self):
        if not self.current_formula:
            # QMessageBox.critical(self, "选择错误", "请先选择一个公式")
            InfoBar.error(
                title='提示',
                content='请先选择一个公式',
                parent=self,
                position=InfoBarPosition.TOP,
                duration=2000
            )
            return

        variables = {}
        for i in range(self.variables_layout.rowCount()):
            label = self.variables_layout.itemAt(i, QFormLayout.LabelRole).widget().text()
            value_edit = self.variables_layout.itemAt(i, QFormLayout.FieldRole).widget()
            try:
                var_name = list(self.current_formula['variables'].keys())[i]
                variables[var_name] = float(value_edit.text())
            except (ValueError, IndexError):
                QMessageBox.critical(self, "输入错误", f"变量 {var_name} 的输入值格式错误，请输入数字")
                return

        required_vars = set(self.current_formula['variables'].keys())
        provided_vars = set(variables.keys())
        if required_vars != provided_vars:
            missing_vars = required_vars - provided_vars
            extra_vars = provided_vars - required_vars
            error_msg = ""
            if missing_vars:
                error_msg += f"缺少变量: {', '.join(missing_vars)}\n"
            if extra_vars:
                error_msg += f"多余的变量: {', '.join(extra_vars)}"
            QMessageBox.critical(self, "变量错误", error_msg.strip())
            return

        # 构建表达式字符串
        expression = self.current_formula['value']
        for var_name, var_value in variables.items():
            expression = expression.replace(var_name, str(var_value))        

        self.expression_edit.setText(expression)  # 显示带入变量值的表达式

        result, error_msg = safe_eval(self.current_formula['value'], variables)
        if error_msg:
            QMessageBox.critical(self, "计算错误", error_msg)
        else:
            self.result_edit.setText(str(result))
    def on_add_edit_clicked(self):
        # 获取当前选中的索引
        current_index = self.tree_view.currentIndex()

        if current_index.isValid():
            # 如果有选中的项，则进行编辑操作
            formula = self.model.get_formula_by_index(current_index)
            if formula:
                # 打开编辑对话框
                self.current_formula_id = formula['id']
                self.formula_name_edit.setText(formula['name'])
                self.formula_expression_edit.setText(formula['value'])
                self.variables_textedit.setText(json.dumps(formula['variables'], indent=4, ensure_ascii=False))
                # 找到对应的分类ID
                index_of_target_value = self.formula_category_combo_box.findData(formula['category_id'])
                # 如果找到了对应的索引，则设置为当前选中项
                if index_of_target_value != -1:
                    self.formula_category_combo_box.setCurrentIndex(index_of_target_value)


                self.right_widget.hide()
                self.formula_widget.show()
                self.del_button.show()
                self.save_button.setText("更新")
            else:
                # 如果选中的是分类，则进行分类编辑
                pass
                
        else:
            # 如果没有选中的项，则进行新增操作
            self.current_formula_id = None
            self.save_button.setText("新增")
            self.right_widget.hide()
            self.formula_widget.show()
            self.del_button.hide()

    def save_formula(self, formula_id=None):
        # 获取输入框中的值
        formula_name = self.formula_name_edit.text()
        formula_expression = self.formula_expression_edit.text()
        # 获取当前选中的分类ID
        formula_category_id = self.formula_category_combo_box.currentData()
        formula_variables = self.variables_textedit.toPlainText()
        # 检查输入是否为空
        if not formula_name or not formula_expression or not formula_variables:
            InfoBar.error(
                title='警告',
                content='请填写完整的公式信息',
                parent=self,
                position=InfoBarPosition.TOP,
                duration=2000
            )
            return

        # 解析变量
        try:
            variables = json.loads(formula_variables)
        except json.JSONDecodeError:
            InfoBar.error(
                title='警告',
                content='变量格式不正确，请使用JSON格式',
                parent=self,
                position=InfoBarPosition.TOP,
                duration=2000
            )
            return

        # 保存公式
        formula={
                'name': formula_name,
                'formula': formula_expression,
                'category_id': formula_category_id,
                'variables': json.dumps(variables)
            }
        if formula_id:
            # 如果有formula_id，则进行更新操作
            formula['id'] = formula_id

        success,msg = self.model.save_or_update_formula(formula)
        
        if success:
            InfoBar.success(
                    title='成功',
                    content='已保存',
                    parent=self,
                    position=InfoBarPosition.TOP,
                    duration=2000
            )
            self.model.load_data()
             # 刷新树视图
            self.tree_view.model().layoutChanged.emit()
            # 关闭窗口
            self.formula_widget.hide()
            self.right_widget.show()
        else:
            InfoBar.error(
                title='警告',
                content=msg,
                parent=self,
                position=InfoBarPosition.TOP,
                duration=2000
            )
            return
       
    
    def on_back_button_clicked(self):
        # 处理返回按钮点击事件的逻辑
        self.formula_widget.hide()
        self.right_widget.show()

    def on_delete_button_clicked(self):
        # 处理删除按钮点击事件的逻辑
        if not self.current_formula:
            InfoBar.error(
                title='警告',
                content='请选择要删除的公式',
                parent=self,
                position=InfoBarPosition.TOP,
                duration=2000
            )
            return
        # 弹出确认对话框
        reply = QMessageBox.question(self, '确认删除', '确定要删除该公式吗？', QMessageBox.Yes | QMessageBox.No)
        if reply == QMessageBox.Yes:
            # 删除当前选中的公式
            success,msg = self.model.delete_formula(self.current_formula_id)
            if success:
                InfoBar.success(
                    title='成功',
                    content='已删除',
                    parent=self,
                    position=InfoBarPosition.TOP,
                    duration=2000
                )
                self.model.load_data()
                # 刷新树视图
                self.tree_view.model().layoutChanged.emit()
                # 关闭窗口
                self.formula_widget.hide()
                self.right_widget.show()
            else:
                InfoBar.error(
                    title='警告',
                    content=msg,
                    parent=self,
                    position=InfoBarPosition.TOP,
                    duration=2000
                )
                return            

